import pygame
import copy
import itertools
from pygame.locals import *

from globals import *
from tile import Tile
from mysprite import MySprite
from snake import BasicSnake

# snake container
class SnakeContainer:
	def __init__(self):
		self.snakes = []
		self.humans = []
		self.bots = []
	
	def getAliveSnakes(self):	
		return filter (lambda snake: snake.isAlive, self.snakes)

	def getAliveBots(self):	
		return filter (lambda bot: bot.isAlive, self.bots)

	def getAliveHumans(self):	
		return filter (lambda human: human.isAlive, self.humans)

	def addSnake(self, snake):
		self.snakes.append(snake)
		if snake.isHuman:
			self.humans.append(snake)
		else:
			self.bots.append(snake)

	def removeSnake(self, snake):
		self.snakes.remove(snake)
		if snake.isHuman:
			self.humans.remove(snake)
		else:
			self.bots.remove(snake)

	def noneAlive(self):
		return not self.getAliveSnakes()

# Snakes after crash are marked as not alive but not removed from 
# snake lists nor grid tiles. Hence, if you iterate over snake list, 
# you might want to check if snakes are alive before working with 
# them. And if you check for collisions you need to discount segments
# of dead snakes that are on the tile. This holds for BasicGrid as
# well as for GameGrid.

# grid logic
class BasicGrid(SnakeContainer):
	# Snakes can be basic or not but this class defines only the basic of them.
	def __init__(self, size):
		SnakeContainer.__init__(self)
		self.gmap = [] # two dimensional array of tiles (lists)
		for x in xrange(size[0]): # build default grid
			if x == 0 or x == size[0] - 1:
				self.gmap.append(
					[Tile((x, y), True) for y in xrange(size[1])])
			else:
				self.gmap.append(
					[Tile((x, 0), True)] +
					[Tile((x, y)) for y in xrange(1,size[1] - 1)] +
					[Tile((x, size[1] - 1), True)])
		self.size = size

	def dump(self):
		for y in xrange(self.size[1]):
			for x in xrange(self.size[0]):
				print self.gmap[x][y].getSymbol(),
			print ""
	
	# check grid for collisions snake&wall or snake&snake&...
	def getCrashedSnakes(self):
		crashedSnakes = []
		for snake in self.snakes:
			if not snake.isAlive: continue
			x, y = snake.segments[0]
			tile = self.gmap[x][y]
			# dead snakes might be present on the tile
			# do not count them.
			aliveSnakes = filter (lambda snake: snake.isAlive, tile.snakes)
			if len(aliveSnakes) > 1 or tile.isWall:
				crashedSnakes.append(snake)
		return set(crashedSnakes) # snake might have hit its own tail

	def update(self):
		# update segments attribute of grid tiles
		for snake in self.snakes:
			if not snake.isAlive: continue
			newHead, oldTail = BasicSnake.update(snake)
			x, y = newHead
			self.gmap[x][y].snakes.append(snake)
			if oldTail == None:
				continue # happens when snake is still growing
			x, y = oldTail
			self.gmap[x][y].snakes.remove(snake)

		# crashed snakes are removed from grid and marked as dead
		# BasicGrid keeps them in its snake lists though, which is
		# important for its child AIGrid.
		crashedSnakes = self.getCrashedSnakes()
		for snake in crashedSnakes:
			# snake.leaveGrid()
			snake.isAlive = False
		return crashedSnakes

# grid also handling sprites of its tiles and snakes
class GameGrid(BasicGrid):
	def __init__(self, size):
		BasicGrid.__init__(self, size)
		self.wallSprites = []
		for x in xrange(len(self.gmap)):
			for y in xrange(len(self.gmap[x])):
				if self.gmap[x][y].isWall:
					self.wallSprites.append(MySprite(WALLCOLOR, (x, y)))
		self.crashedSnakes = []

	def isEmpty(self):
		return self.noneAlive() and not self.crashedSnakes

	def update(self, sprites):
		# update snakes attribute of grid tiles
		for snake in self.snakes:
			if not snake.isAlive: continue
			newHead, oldTail = snake.update(sprites)
			x, y = newHead
			self.gmap[x][y].snakes.append(snake)
			if oldTail == None:
				continue # happens when snake is still growing
			x, y = oldTail
			self.gmap[x][y].snakes.remove(snake)

		# simple animation after crashing
		for snake in self.crashedSnakes:
			# may set snake.inHeaven flag
			snake.crashAnimation()
		# remove snakes that are already in heaven
		self.crashedSnakes[:] = [s for s in self.crashedSnakes if not s.inHeaven]

		# drunk or owned snakes
		for snake in self.getCrashedSnakes():
			snake.isAlive = False
			snake.leaveGrid(sprites)
			self.crashedSnakes.append(snake)
